"""
A simple version of Proximal Policy Optimization (PPO) using single thread.
Based on:
    1. Emergence of Locomotion Behaviours in Rich Environments (Google Deepmind): [https://arxiv.org/abs/1707.02286]
    2. Proximal Policy Optimization Algorithms (OpenAI): [https://arxiv.org/abs/1707.06347]
View more on my tutorial website: https://morvanzhou.github.io/tutorials
Dependencies:
tensorflow r1.2
gym 0.9.2
"""
# Reference:
# https://morvanzhou.github.io/tutorials/machine-learning/reinforcement-learning/6-4-DPPO/

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
import gym

EP_MAX = 1000
Steps = 200
GAMMA = 0.9

A_LR = 0.0001
C_LR = 0.0002

BATCH = 32
A_UPDATE_STEPS = 10
C_UPDATE_STEPS = 10
S_DIM, A_DIM = 3, 1

# clipped for default
METHOD = [
    dict(name='kl_pen', kl_target=0.01, lam=0.5),   # KL penalty
    dict(name='clip', epsilon=0.2),                 # Clipped surrogate objective, find this is better
][1]        


class PPO(object):
    def __init__(self):
        self.sess = tf.Session()        
        # <tf.Tensor 'state:0' shape=(?, 3) dtype=float32>
        self.tfs = tf.placeholder(tf.float32, [None, S_DIM], 'state')
        # critic
        # 定义critic计算优势函数,状态价值定义了一个全联接神经网络来得到
        # tf.layers.dense()
        # 该层实现了操作：outputs = activation(inputs * kernel + bias),
        # 其中activation是作为activation参数传递的激活函数(如果不是None),是由层创建的权重矩阵,kernel是由层创建的权重矩阵,并且bias是由层创建的偏差向量(只有use_bias为True时).
        with tf.variable_scope('critic'):
            # <tf.Tensor 'critic/dense/Relu:0' shape=(?, 100) dtype=float32>
            l1 = tf.layers.dense(self.tfs, 100, tf.nn.relu)
            # 优势函数，即Qπ(st,at) - Vπ(st)。其中Qπ(st,at)可以使用从当前状态开始到eposide结束的奖励折现和得到，
            # Vπ(st)可以通过一个critic来计算得到
            self.v = tf.layers.dense(l1, 1) # <tf.Tensor 'critic/dense_1/BiasAdd:0' shape=(?, 1) dtype=float32>
            self.tfdc_r = tf.placeholder(tf.float32, [None, 1], 'discounted_r')
            self.advantage = self.tfdc_r - self.v
            # tf.reduce_mean 函数用于计算张量tensor沿着指定的数轴（tensor的某一维度）上的的平均值，主要用作降维或者计算tensor（图像）的平均值
            self.closs = tf.reduce_mean(tf.square(self.advantage))
            # 更新 Critic 的时候是根据 刚刚计算的 discounted_r 和自己分析出来的 state_value 
            # 这两者的差 (advantage),然后最小化这个差值.类似于TD_error
            # C_LR: critic_learning_rate
            # 我们更新 Critic 的时候是根据 刚刚计算的 discounted_r 和自己分析出来的 state value 这两者的差 (advantage). 然后最小化这个差值
            self.ctrain_op = tf.train.AdamOptimizer(C_LR).minimize(self.closs)

        # actor
        # PPO里采用的是off-policy的策略，需要有一个单独的网络来收集数据，并用于策略的更新，
        # 同DQN的策略一样，我们定义了一个单独的网络，这个网络的参数是每隔一段时间由我们真正的Actor的参数复制过去
        pi, pi_params = self._build_anet('pi', trainable=True)  # actor's policy
        oldpi, oldpi_params = self._build_anet('oldpi', trainable=False) # old actor's policy

        # choosing action
        with tf.variable_scope('sample_action'):
            self.sample_op = tf.squeeze(pi.sample(1), axis=0)  # Generate samples of the specified shape
        # 每次要进行 PPO 更新 Actor 和 Critic 的时候, 我们有需要将 pi 的参数复制给 oldpi
        with tf.variable_scope('update_oldpi'):
            self.update_oldpi_op = [oldp.assign(p) for p, oldp in zip(pi_params, oldpi_params)]

        self.tfa = tf.placeholder(tf.float32, [None, A_DIM], 'action')
        self.tfadv = tf.placeholder(tf.float32, [None, 1], 'advantage')

        # 损失计算,采用KL散度或者裁切的方式，损失计算方式不同
        with tf.variable_scope('loss'):
            with tf.variable_scope('surrogate'):
                # ratio = tf.exp(pi.log_prob(self.tfa) - oldpi.log_prob(self.tfa))
                ratio = pi.prob(self.tfa) / oldpi.prob(self.tfa) # Probability density/mass function
                surr = ratio * self.tfadv
            if METHOD['name'] == 'kl_pen':  # 如果用 KL penatily
                self.tflam = tf.placeholder(tf.float32, None, 'lambda')
                kl = tf.distributions.kl_divergence(oldpi, pi)
                self.kl_mean = tf.reduce_mean(kl)
                self.aloss = -(tf.reduce_mean(surr - self.tflam * kl))
            else:   # clipping method, find this is better    # 如果用 clipping 的方式
                self.aloss = -tf.reduce_mean(tf.minimum(
                    surr,
                    tf.clip_by_value(ratio, 1.-METHOD['epsilon'], 1.+METHOD['epsilon'])*self.tfadv))

        with tf.variable_scope('atrain'):
            self.atrain_op = tf.train.AdamOptimizer(A_LR).minimize(self.aloss)

        tf.summary.FileWriter("log/", self.sess.graph)

        self.sess.run(tf.global_variables_initializer())

    # 更新 PPO
    def update(self, s, a, r):

        self.sess.run(self.update_oldpi_op)
        # 得到advantage value
        adv = self.sess.run(self.advantage, {self.tfs: s, self.tfdc_r: r})
        # adv = (adv - adv.mean())/(adv.std()+1e-6)     # sometimes helpful

        # update actor
        if METHOD['name'] == 'kl_pen':
            for _ in range(A_UPDATE_STEPS):
                _, kl = self.sess.run(
                    [self.atrain_op, self.kl_mean],
                    {self.tfs: s, self.tfa: a, self.tfadv: adv, self.tflam: METHOD['lam']})
                if kl > 4*METHOD['kl_target']:  # this in in google's paper
                    break
            if kl < METHOD['kl_target'] / 1.5:  # adaptive lambda, this is in OpenAI's paper
                METHOD['lam'] /= 2
            elif kl > METHOD['kl_target'] * 1.5:
                METHOD['lam'] *= 2
            METHOD['lam'] = np.clip(METHOD['lam'], 1e-4, 10)    # sometimes explode, this clipping is my solution
        else:   # clipping method, find this is better (OpenAI's paper)
            [self.sess.run(self.atrain_op, {self.tfs: s, self.tfa: a, self.tfadv: adv}) for _ in range(A_UPDATE_STEPS)]
        # 这个 update 的步骤里, 我们用 for loop 更新了很多遍 Actor 和 Critic, 
        # 在 loop 之前, pi 和 old pi 是一样的, 每次 loop 的之后, pi 会变动, 而 old pi 不变, 
        # 这样这个 surrogate surrogate 就会开始变动了. 这就是 PPO 的精辟

        # update critic  # 更新 Critic 的时候, 他们是一样的
        # self.tfs: s  状态   # self.tfdc_r: r 目标状态价值函数
        [self.sess.run(self.ctrain_op, {self.tfs: s, self.tfdc_r: r}) for _ in range(C_UPDATE_STEPS)]

    # 我们的Actor网络输出一个均值和方差，并返回一个由该均值和方差得到的正态分布，动作基于此正态分布进行采样
    def _build_anet(self, name, trainable):
        with tf.variable_scope(name):
            l1 = tf.layers.dense(self.tfs, 100, tf.nn.relu, trainable=trainable)
            mu = 2 * tf.layers.dense(l1, A_DIM, tf.nn.tanh, trainable=trainable)
            sigma = tf.layers.dense(l1, A_DIM, tf.nn.softplus, trainable=trainable)
            norm_dist = tf.distributions.Normal(loc=mu, scale=sigma)  # 一个正态分布
        params = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope=name)
        return norm_dist, params

    # 选动作
    def choose_action(self, s):
        s = s[np.newaxis, :]
        a = self.sess.run(self.sample_op, {self.tfs: s})[0]
        return np.clip(a, -2, 2)

    # 算 state value
    def get_v(self, s):
        if s.ndim < 2: s = s[np.newaxis, :]
        return self.sess.run(self.v, {self.tfs: s})[0, 0]

env = gym.make('Pendulum-v0').unwrapped
ppo = PPO()
all_ep_r = []

for ep in range(EP_MAX):
    s = env.reset()
    buffer_s, buffer_a, buffer_r = [], [], []
    ep_r = 0
    for t in range(Steps):    # in one episode
        env.render()
        a = ppo.choose_action(s)  # 根据一个正态分布，选择一个action
        s_, r, done, _ = env.step(a)
        buffer_s.append(s)
        buffer_a.append(a)
        buffer_r.append((r+8)/8)    # normalize reward # 发现有帮助
        s = s_
        ep_r += r

        # update ppo  # 如果 buffer 收集一个 batch 了或者 episode 完了
        if (t+1) % BATCH == 0 or t == Steps-1:
            v_s_ = ppo.get_v(s_)  # get_v()是一个得到状态价值的函数。使用v(s) = r + gamma * v(s+1)不断进行循环
            discounted_r = []
            # 计算 discounted reward
            for r in buffer_r[::-1]:
                v_s_ = r + GAMMA * v_s_
                discounted_r.append(v_s_)
            discounted_r.reverse()

            bs, ba, br = np.vstack(buffer_s), np.vstack(buffer_a), np.array(discounted_r)[:, np.newaxis]
            # 清空 buffer
            buffer_s, buffer_a, buffer_r = [], [], []
            ppo.update(bs, ba, br)  # 更新 PPO

    if ep == 0: all_ep_r.append(ep_r)
    else: all_ep_r.append(all_ep_r[-1]*0.9 + ep_r*0.1)
    print(
        'Ep: %i' % ep,
        "|Ep_r: %i" % ep_r,
        ("|Lam: %.4f" % METHOD['lam']) if METHOD['name'] == 'kl_pen' else '',
    )

plt.plot(np.arange(len(all_ep_r)), all_ep_r)
plt.xlabel('Episode');plt.ylabel('Moving averaged episode reward');plt.show()